package com.tuniu.agents.config;

import reactor.core.publisher.Flux;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * General Stream Message Processor
 * Provides core stream processing functionality
 */
public class StreamCommonProcessor implements StreamProcessor {

    // Regular expression cache to improve performance
    private static final ConcurrentHashMap<String, Pattern> PATTERN_CACHE = new ConcurrentHashMap<>();

    /**
     * Process captured content; subclasses can override this method for custom processing
     */
    @Override
    public String processCapturedContent(String content) {
        return content;
    }

    /**
     * Process the stream using a custom extraction logic
     * Uses the handle operator to avoid subscribing in a non-blocking context
     *
     * @param flux Input stream
     * @param extractor Extractor function
     * @param emitOriginal Whether to emit the original chunk
     * @return Processed stream
     */
    public Flux<String> process(Flux<String> flux, Function<String, ExtractResult> extractor, boolean emitOriginal) {
        return Flux.defer(() -> {
            StringBuilder buffer = new StringBuilder(4096);

            return flux.flatMap(chunk -> {
                // Decide whether to emit the original chunk based on the parameter
                Flux<String> originalChunk = emitOriginal ? Flux.just(chunk) : Flux.empty();

                buffer.append(chunk);
                String content = buffer.toString();

                // Use the extractor to process the content
                ExtractResult result = extractor.apply(content);

                // If there is extracted content, add it to the result stream
                Flux<String> extractedContent = Flux.empty();
                if (result.hasExtracted()) {
                    extractedContent = Flux.fromIterable(result.getExtractedContents())
                            .map(this::processCapturedContent);

                    // Update buffer to keep only the unprocessed part
                    if (result.getProcessedLength() > 0) {
                        buffer.delete(0, result.getProcessedLength());
                    }
                }

                // Merge original chunk and extracted content
                return Flux.concat(originalChunk, extractedContent);
            });
        });
    }

    /**
     * Process the stream using a custom extraction logic, default to returning the original stream
     *
     * @param flux Input stream
     * @param extractor Extractor function
     * @return Processed stream
     */
    @Override
    public Flux<String> process(Flux<String> flux, Function<String, ExtractResult> extractor) {
        return process(flux, extractor, true);
    }

    /**
     * Process the stream using start and end markers
     */
    public Flux<String> processWithStartAndEnd(Flux<String> flux, String start, String end, boolean emitOriginal) {
        return process(flux, content -> extractByStartAndEnd(content, start, end), emitOriginal);
    }

    /**
     * Process the stream using start and end markers, default to returning the original stream
     */
    public Flux<String> processWithStartAndEnd(Flux<String> flux, String start, String end) {
        return processWithStartAndEnd(flux, start, end, true);
    }

    /**
     * Process based on start marker and fixed character count
     */
    public Flux<String> processWithCharCount(Flux<String> flux, String start, int charCount, boolean emitOriginal) {
        return process(flux, content -> extractByCharCount(content, start, charCount), emitOriginal);
    }

    /**
     * Process based on start marker and fixed character count, default to returning the original stream
     */
    public Flux<String> processWithCharCount(Flux<String> flux, String start, int charCount) {
        return processWithCharCount(flux, start, charCount, true);
    }

    /**
     * Process the stream using regular expressions
     */
    public Flux<String> processWithRegex(Flux<String> flux, String regex, int groupIndex, boolean emitOriginal) {
        return processWithRegex(flux, regex, groupIndex, null, emitOriginal);
    }

    /**
     * Process the stream using regular expressions, default to returning the original stream
     */
    public Flux<String> processWithRegex(Flux<String> flux, String regex, int groupIndex) {
        return processWithRegex(flux, regex, groupIndex, true);
    }

    /**
     * Process the stream using regular expressions, supports adding prefixes and suffixes
     */
    public Flux<String> processWithRegex(Flux<String> flux, String regex, int groupIndex, ContentWrapper wrapper, boolean emitOriginal) {
        return process(flux, content -> {
            ExtractResult result = extractByRegex(content, regex, groupIndex);
            if (wrapper != null && result.hasExtracted()) {
                List<String> wrappedContents = result.getExtractedContents().stream()
                        .map(wrapper::wrap)
                        .toList();
                result.setExtractedContents(wrappedContents);
            }
            return result;
        }, emitOriginal);
    }

    /**
     * Process the stream using regular expressions, supports adding prefixes and suffixes, default to returning the original stream
     */
    public Flux<String> processWithRegex(Flux<String> flux, String regex, int groupIndex, ContentWrapper wrapper) {
        return processWithRegex(flux, regex, groupIndex, wrapper, true);
    }

    /**
     * Extract content based on start and end markers
     */
    public static ExtractResult extractByStartAndEnd(String content, String start, String end) {
        ExtractResult result = new ExtractResult();

        if (content == null || content.length() < start.length() + end.length()) {
            return result;
        }

        int startIndex;
        int searchPos = 0;
        int contentLength = content.length();
        int startLength = start.length();
        int endLength = end.length();

        int estimatedMatches = Math.max(1, contentLength / (startLength + endLength + 10));
        result.ensureCapacity(estimatedMatches);

        while (searchPos < contentLength && (startIndex = content.indexOf(start, searchPos)) != -1) {
            int contentStartPos = startIndex + startLength;
            if (contentStartPos >= contentLength) {
                break;
            }

            int endIndex = content.indexOf(end, contentStartPos);
            if (endIndex == -1) {
                break;
            }

            String extracted = content.substring(contentStartPos, endIndex);
            result.addExtractedContent(extracted);
            searchPos = endIndex + endLength;
        }

        result.setProcessedLength(searchPos);
        return result;
    }

    /**
     * Extract content based on character count
     */
    public static ExtractResult extractByCharCount(String content, String start, int charCount) {
        ExtractResult result = new ExtractResult();

        if (content == null || content.length() < start.length() + 1) {
            return result;
        }

        int startIndex;
        int searchPos = 0;
        int contentLength = content.length();
        int startLength = start.length();

        int estimatedMatches = Math.max(1, contentLength / (startLength + charCount));
        result.ensureCapacity(estimatedMatches);

        while (searchPos < contentLength && (startIndex = content.indexOf(start, searchPos)) != -1) {
            int contentStartPos = startIndex + startLength;

            if (contentStartPos + charCount <= contentLength) {
                String extracted = content.substring(contentStartPos, contentStartPos + charCount);
                result.addExtractedContent(extracted);
                searchPos = contentStartPos + charCount;
            } else {
                break;
            }
        }

        result.setProcessedLength(searchPos);
        return result;
    }

    /**
     * Extract content based on regular expressions
     */
    public static ExtractResult extractByRegex(String content, String regex, int groupIndex) {
        ExtractResult result = new ExtractResult();

        if (content == null || content.isEmpty()) {
            return result;
        }

        Pattern pattern = PATTERN_CACHE.computeIfAbsent(regex, Pattern::compile);
        Matcher matcher = pattern.matcher(content);

        int lastEnd = 0;

        int estimatedMatches = Math.max(1, content.length() / 20);
        result.ensureCapacity(estimatedMatches);

        while (matcher.find()) {
            if (groupIndex >= 0 && groupIndex <= matcher.groupCount()) {
                String extracted = matcher.group(groupIndex);
                if (extracted != null) {
                    result.addExtractedContent(extracted);
                }
            }
            lastEnd = matcher.end();
        }

        result.setProcessedLength(lastEnd);
        return result;
    }
}